{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Image-Based Wavefront Sensing\n",
    "\n",
    "In this example, we will show how ot implement image-based wavefront sensing based on prysm.  We will use the library both to synthesize the truth data, and as the embedded modeling tool used as part of the optimization routine.\n",
    "\n",
    "We begin by importing a few classes and functions from prysm and other modules and writing some functions to generate data from zernike coefficients,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import partial\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "from prysm import NollZernike, PSF\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "%matplotlib inline\n",
    "plt.style.use('bmh')\n",
    "\n",
    "def bake_in_defocus(zernikes, defocus_values):\n",
    "    return [{**zernikes, **dict(Z4=defocus)} for defocus in defocus_values]\n",
    "\n",
    "def zerns_idxs_to_dict(zernike_coefs, indices):\n",
    "    return {k:v for k, v in zip(indices, zernike_coefs)}\n",
    "\n",
    "# a few extra variables are passed in and aren't used,\n",
    "# but we're aiming for brevity in this example, not pristine code\n",
    "# so we just collect them with kwargs and don't use them\n",
    "def zernikes_to_pupil(zerns, epd, wvl, samples, efl=None, Q=None, target='fcn'):\n",
    "    return NollZernike(**zerns,  # coefficients\n",
    "                       dia=epd, wavelength=wvl,  # physical parameters\n",
    "                       norm=True, opd_unit='waves',  # units and normalization\n",
    "                       mask='circle', mask_target=target,  # geometry\n",
    "                       samples=samples)  # sampling\n",
    "\n",
    "def zernikes_to_psfs(sets_of_zernikes, efl, epd, Q, wvl, samples):\n",
    "    psfs = []\n",
    "    for zerns in sets_of_zernikes:\n",
    "        pupil = zernikes_to_pupil(zerns, epd, wvl, samples)\n",
    "        psf = PSF.from_pupil(pupil, Q=Q, efl=efl)\n",
    "        psfs.append(psf)\n",
    "    return psfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`zernikes_to_psfs` is quite terse, but has a lot going on.  `zerns` is a dictionary that contains key-value pairs of (Noll) Zernike indexes and their coefficients.  Unpacking it allows us use arbitrary combinations of Zernike terms within the algorithm.  In `zernikes_to_pupil,` `mask_target` is used to avoid masking the phase during optimization, improving performance.  We make explicit that the pupil is circular with `mask`.  It will be used both inside the optimizer and to synthesize data.  We will treat EFL, EPD, Q, and $\\lambda$ as known.\n",
    "\n",
    "Next, we synthesize the truth data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# wavefront data\n",
    "defocus_values_waves_rms = [-0.65, 0, 0.76]  # noninteger values with nothing special about them\n",
    "reference_zernikes = dict(Z5=0.012, Z6=0.023, Z7=0.034, Z8=0.045, Z9=0.056, Z10=0.067, Z11=0.12)\n",
    "wavefront_coefs = bake_in_defocus(reference_zernikes, defocus_values_waves_rms)\n",
    "\n",
    "# system parameters\n",
    "efl = 1500\n",
    "epd = 150  # F/10, e.g. an unobscured telescope with 15 cm aperture\n",
    "wvl = 0.55 # monochromatic visible system\n",
    "Q = 2.13   # minorly oversampled\n",
    "samples = 128 # this is a reasonable number for small OPD and uncomplicated aperture geometry\n",
    "\n",
    "ztp_kwargs = dict(efl=efl, epd=epd, Q=Q, wvl=wvl, samples=samples)\n",
    "truth_psfs = zernikes_to_psfs(wavefront_coefs, **ztp_kwargs)\n",
    "\n",
    "def rowplot_psfs(psfs):\n",
    "    fig, axs = plt.subplots(ncols=3, figsize=(10,4))\n",
    "    for psf, ax in zip(psfs, axs):\n",
    "        psf.plot2d(fig=fig, ax=ax, axlim=125, power=2)\n",
    "        ax.grid(False)\n",
    "\n",
    "    fig.tight_layout()\n",
    "    return fig, axs\n",
    "\n",
    "rowplot_psfs(truth_psfs);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The goal of image-based wavefront sensing is to take these PSFs and use them to estimate the wavefront that generated them (above, `reference_zernikes`) without prior knowledge.  We do this by constructing a nonlinear optimization problem and asking the computer to determine the wavefront coefficients for us.  To that end, we formulate a cost function which calculates the mean square error between the data and the truth,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optfcn(zernike_coefs, indices, ztp_kwargs, defocuses, truth_psfs):\n",
    "    base_wavefront_coefs = zerns_idxs_to_dict(zernike_coefs, indices)\n",
    "    wavefront_coefs = bake_in_defocus(base_wavefront_coefs, defocuses)\n",
    "    psfs = zernikes_to_psfs(wavefront_coefs, **ztp_kwargs)\n",
    "    \n",
    "    # t = truth, m = model.  sum(^2) = mean square error\n",
    "    diffs = [((t.data - m.data)**2).sum() for t, m in zip(truth_psfs, psfs)]\n",
    "    \n",
    "    # normalize by N,\n",
    "    # psf.size = number of pixels\n",
    "    diff = sum(diffs) / (len(psfs) * psfs[0].size)\n",
    "    return diff\n",
    "\n",
    "coefs = [f'Z{n}' for n in [5, 6, 7, 8, 9, 10, 11]]\n",
    "optfcn_used = partial(optfcn,\n",
    "                      indices=coefs,\n",
    "                      ztp_kwargs=ztp_kwargs,\n",
    "                      defocuses=defocus_values_waves_rms,\n",
    "                      truth_psfs=truth_psfs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This function computes the mean square error between the data in our model PSFs and the truth.  Let's check that it returns zero for the truth:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optfcn_used(reference_zernikes.values())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This (probably) means we didn't make a mistake.  In your own program, this should be verified more rigorously.  To get a sense for execution time, how quickly can we evaluate it?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit optfcn_used(reference_zernikes.values())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "On the computer this document was written on, the time is 5.5ms per run.  So we should expect much less than a second per iteration of the optimizer; this problem is not so large that running it requires a supercomputer.\n",
    "\n",
    "All that is left to do is ask the optimizer kindly for the true wavefront, given a guess.  A general guess might be all zeros -- a pupil with no wavefront error.  This does not assume any prior knowledge.  The L-BFGS-B minimizer tends to perform the best for this problem, from experience."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opt_result = minimize(optfcn_used, tuple([0] * len(coefs)), method='L-BFGS-B')\n",
    "opt_result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Casting the `x0` variable to a tuple makes it immutable, which will help when you try to debug one of these problems, but doesn't matter here.\n",
    "\n",
    "How did the optimizer do?  Well, let's compare the PSFs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# this is copy pasted from optfcn, but we aren't trying to make the cleanest code right now.\n",
    "base_wavefront_coefs = zerns_idxs_to_dict(opt_result.x, coefs)\n",
    "wavefront_coefs = bake_in_defocus(base_wavefront_coefs, defocus_values_waves_rms)\n",
    "retrieved_psfs = zernikes_to_psfs(wavefront_coefs, **ztp_kwargs)\n",
    "\n",
    "print('true PSFs')\n",
    "rowplot_psfs(truth_psfs);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('estimated PSFs')\n",
    "rowplot_psfs(retrieved_psfs);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Looks like a good match.  How about the wavefronts?  We have to change the mask target here to the phase visualizes the way we expect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "true_wavefront = zernikes_to_pupil(reference_zernikes, **ztp_kwargs, target='both')\n",
    "retrieved_wavefront = zernikes_to_pupil(wavefront_coefs[1], **ztp_kwargs, target='both')\n",
    "elementwise_difference = true_wavefront - retrieved_wavefront\n",
    "\n",
    "fig, axs = plt.subplots(ncols=3, figsize=(14,8))\n",
    "true_wavefront.plot2d(fig=fig, ax=axs[0])\n",
    "retrieved_wavefront.plot2d(fig=fig, ax=axs[1])\n",
    "elementwise_difference.plot2d(fig=fig, ax=axs[2])\n",
    "\n",
    "\n",
    "for ax in axs:\n",
    "    ax.grid(False)\n",
    "\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The optimization function ideally lets us predict the RMS error of the estimate.  Let's compare,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.sqrt(opt_result.fun), elementwise_difference.rms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "They differ by about a factor of ten, which is unfortunate.  The RMS error is 2 thousandths of a wave; a pretty good result.  What's the peak error?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "elementwise_difference.pv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "about 1 1/2 hundredth of a wave, or ~7 nanometers.  Not bad.  Competitive with interferometers, anyway."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More advanced tasks are left to the reader, such as:\n",
    "\n",
    "* handling of noise\n",
    "\n",
    "* performance optimization\n",
    "\n",
    "* improved flexibility w.r.t units, pupil shapes, etc\n",
    "\n",
    "* estimation of Q or other system parameters\n",
    "\n",
    "* inclusion of focal plane effects\n",
    "\n",
    "* extension to N PSFs instead of 3\n",
    "\n",
    "* use of more terms\n",
    "\n",
    "* use of other types of phase diversity instead of focus\n",
    "\n",
    "* formulation of a cost function better connected to the error in the wavefront estimate\n",
    "\n",
    "* tracking of the optimizer to better understand the course of optimization\n",
    "\n",
    "* use of other data, such as [MTF](https://static1.squarespace.com/static/578d10066a4963fd85e0aa32/t/5af25e61aa4a99ed9ecfa39c/1525833475522/bdd_ug_thesis_10.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}